#!/usr/bin/python2
# -*- coding: utf-8 -*-
"""
-------------------------------------------------
   File Name：     coolKprobe
   Description :
   Author :       chuyansz
   date：          2021/7/4
-------------------------------------------------
   Change Activity:
                   2021/7/4:
-------------------------------------------------
"""
__author__ = 'chuyansz'

import re
import signal
import sys
import os
import socket
import shlex
import argparse
import struct
from subprocess import PIPE, Popen
from threading import Thread
from Queue import Queue, Empty
ON_POSIX = 'posix' in sys.builtin_module_names

def getValueFromStr(s):
    if s.startswith('0x') or s.startswith('0X'):
        return int(s, 16)
    else:
        return int(s)

def headTrans(head, value):
    type = head.split('_', 1)[0]
    if type == "ip":
        v = int(value, 16)
        value = socket.inet_ntoa(struct.pack('>I', socket.htonl(v)))
    elif type == "B16":
        v = getValueFromStr(value)
        v = struct.unpack('H', struct.pack('>H', v))
        value = "%x" % v
    elif type == "B32":
        v = getValueFromStr(value)
        v = struct.unpack('I', struct.pack('>I', v))
        value = "%x" % v
    elif type == "B64":
        v = getValueFromStr(value)
        v = struct.unpack('Q', struct.pack('>Q', v))
        value = "%x" % v
    elif type == "b16":
        v = getValueFromStr(value)
        v = struct.unpack('H', struct.pack('>H', v))
        value = "%d" % v
    elif type == "b32":
        v = getValueFromStr(value)
        v = struct.unpack('I', struct.pack('>I', v))
        value = "%d" % v
    elif type == "b64":
        v = getValueFromStr(value)
        v = struct.unpack('Q', struct.pack('>Q', v))
        value = "%d" % v
    return "%s=%s" % (head, value)

def invHeadTrans(head, value):  #为 filter 翻转
    type = head.split('_', 1)[0]
    if type == "ip":
        v = struct.unpack('I',socket.inet_aton(value))[0]
        value = "0x%x" % v
    elif type == "B16" or type == "b16":
        v = getValueFromStr(value)
        v = struct.unpack('H', struct.pack('>H', v))
        value = "%x" % v
    elif type == "B32" or type == 'b32':
        v = getValueFromStr(value)
        v = struct.unpack('I', struct.pack('>I', v))
        value = "%x" % v
    elif type == "B64" or type == 'b64':
        v = getValueFromStr(value)
        v = struct.unpack('Q', struct.pack('>Q', v))
        value = "%x" % v
    return head, value

class RootRequiredException(Exception):
    def __init__(self, msg=""):
        super(RootRequiredException, self).__init__(msg)

class FileNotExistException(Exception):
    def __init__(self, msg=""):
        super(FileNotExistException, self).__init__(msg)

class FileNotEmptyException(Exception):
    def __init__(self, msg=""):
        super(FileNotEmptyException, self).__init__(msg)

class InvalidArgsException(Exception):
    def __init__(self, msg=""):
        super(InvalidArgsException, self).__init__(msg)

class CasyncPipe(Thread):
    def __init__(self, f, func):
        if not os.path.exists(f):
            FileNotExistException("%s is not exist." % f)
        self.__callBack = func
        super(CasyncPipe, self).__init__()
        self.daemon = True  # thread dies with the program
        self.__pipe = open(f, 'r')
        self.__loop = True
        self.start()

    def run(self):
        while self.__loop:
            line = self.__pipe.readline().strip()
            self.__callBack(line)

    def terminate(self):
        self.__loop = False
        self.join(1)

class CexecCmd(object):
    def __init__(self):
        pass

    def cmd(self, cmds):
        p = Popen(shlex.split(cmds), stdout=PIPE)
        return p.stdout.read().strip()

    def system(self, cmds):
        cmds = cmds.replace('\0', '').strip()
        return os.popen(cmds)

class CasyncCmdQue(Thread):
    def __init__(self, cmd):
        super(CasyncCmdQue, self).__init__()
        self.daemon = True  # thread dies with the program

        self.__q = Queue()
        self.__p = Popen(shlex.split(cmd), stdout=PIPE, stdin=PIPE, close_fds=ON_POSIX)
        self.__stopFlag = False

        self.start()

    def run(self):
        out = self.__p.stdout
        for line in iter(out.readline, b''):
            self.__q.put(line)
            if self.__stopFlag:
                break
        out.close()

    def readLines(self, timeout=0.1):
        lines = []
        while True:
            try:
                line = self.__q.get(timeout=timeout)
            except Empty:
                break
            else:  # got line
                lines.append(line)
        return lines

    def readLine(self, timeout=0.1):
        try:
            line = self.__q.get(timeout=timeout)
        except Empty:
            return None
        else:  # got line
            return line

    def read(self, timeout=0.1):
        return "".join(self.readLines(timeout))

    def write(self, cmd):
        try:
            return self.__p.stdin.write(cmd)
        except IOError:
            return -1

    def writeLine(self, cmd):
        try:
            return self.__p.stdin.write(cmd + '\n')
        except IOError:
            return -1

    def terminate(self):
        self.__stopFlag = True
        self.__p.terminate()
        self.join(1)
        return self.__p.wait()

    def sendSignal(self, sig):
        self.__p.send_signal(sig)

class ftrace(object):
    def __init__(self, show=False):
        super(ftrace, self).__init__()
        self._c = CexecCmd()
        if show:
            self.baseDir = ""
        else:
            self.__checkRoot()
            self.baseDir = self.__getMountDir()
        self.pipe = None
        self._stopHook = []
        self._single = True
        self.__ps = []

    def __getMountDirStr(self):
        cmd = "mount"
        lines = self._c.cmd(cmd)
        for l in lines.split('\n'):
            if "type debugfs" in l:
                return l
        return None

    def __checkRoot(self):
        cmd = 'whoami'
        line = self._c.cmd(cmd).strip()
        if line != "root":
            raise RootRequiredException('this app need run as root')

    def __getMountDir(self):
        s = self.__getMountDirStr()
        if s is None:
            return None
        else:
            return s.split(' ')[2]

    def tracePre(self, buffSize=2048):
        pBuffersize = self.baseDir + "/tracing/buffer_size_kb"
        self._echoPath(pBuffersize, "%d" % buffSize)
        pTrace = self.baseDir + "/tracing/trace"
        self._echoPath(pTrace)

    def _echoPath(self, path, value=""):
        cmd = "echo %s > %s" % (value, path)
        print(cmd)
        return self._c.system(cmd)

    def _echoDPath(self, path, value=""):
        cmd = "echo %s >> %s" % (value, path)
        print(cmd)
        return self._c.system(cmd)

    def procLine(self, line):
        print(line)
        return 0

    def __stopTracing(self):
        pOn = self.baseDir + "/tracing/tracing_on"
        self._echoPath(pOn, "0")

    def _start(self):
        pOn = self.baseDir + "/tracing/tracing_on"
        self._echoPath(pOn, "1")
        self._stopHook.insert(0, self.__stopTracing)
        signal.signal(signal.SIGINT, self.signalHandler)

    def start(self):
        self._single = True
        self._start()
        pipe = self.baseDir + "/tracing/trace_pipe"
        self.pipe = CasyncPipe(pipe, self.procLine)

    def signalHandler(self, signalNumber, frame):
        if signalNumber == signal.SIGINT:
            self.stop()

    def stop(self):
        if self._single:
            self.pipe.terminate()
        else:
            for p in self.__ps:
                print("stop %d" % p.pid)
                p.join()
        for hook in self._stopHook:
            hook()

    def loop(self):
        signal.pause()

class CpaserGdbArgs(object):
    def __init__(self, vmlinuxPath=None):
        self._checkGdbExist()
        if vmlinuxPath is None:
            vmlinuxPath = self.__getVmlinuxPath()

        if not os.path.exists(vmlinuxPath):
            raise FileNotExistException("vmlinux %s is not found." % (vmlinuxPath))

        self.__re1 = re.compile("(?<=<)[^<>]+(?=>)")
        self.__reIndex = re.compile("(?<=\\[)[^\\[\\]]+(?=\\])")
        cmd = 'gdb %s' % vmlinuxPath
        self.__aCmd = CasyncCmdQue(cmd)
        self.__aCmd.read(1)
        self.__aCmd.writeLine('set pagination off')
        self.__aCmd.read()

    def __del__(self):
        self.__aCmd.terminate()

    def _checkGdbExist(self):
        cmd = 'which gdb'
        c = CexecCmd()
        line = c.cmd(cmd).strip()
        if line == "":
            raise FileNotFountException("gdb is not install")

    def __getVmlinuxPath(self):
        c = CexecCmd()
        name = c.cmd('uname -r')
        return "/usr/lib/debug/usr/lib/modules/" + name + "/vmlinux"

    def showTypeSize(self, _type):
        cmd = "p sizeof(%s)" % _type
        self.__aCmd.writeLine(cmd)
        readStr = self.__aCmd.readLine(1).strip()
        nStr = readStr.split('=')[1].strip()
        return int(nStr)

    def showStruct(self, structName):
        self.__aCmd.writeLine("ptype %s" % structName)
        return self.__aCmd.read(0.5)

    def __checkIndex(self, i, l, member):
        finds = self.__reIndex.findall(l)
        if len(finds) == 0: #  empty, means []， do not check
            return
        else:
            size = int(finds[0])
            if i >= size:
                raise InvalidArgsException("member %s index %d, max %d, overflow." % (member, i, size))

    def memberTypes(self, member, structName):
        structName = structName.replace("*", "")
        lines = self.showStruct(structName).split('\n')
        if '[' in member:
            m = member.split('[')[0]
            n = int(self.__reIndex.findall(member)[0])
            for l in lines:
                l = l.strip()
                find1 = " " + m + '['
                if find1 in l:
                    self.__checkIndex(n, l, member)
                    return l.rsplit(' ', 1)[0]
                find2 = "*" + m + '['
                if find2 in l:
                    self.__checkIndex(n, l, member)
                    return l.rsplit(' ', 1)[0] + '*'
            raise InvalidArgsException("%s is not in %s" % (member, structName))
        else:
            for l in lines:
                l = l.strip()
                find1 = " " + member + ";"
                if l.endswith(find1):
                    return l[:-len(find1)].strip()
                find2 = " *" + member + ";"
                if l.endswith(find2):
                    return l[:-len(find2)].strip() + " *"
                find3 = "char %s[" % member  # char a[16]
                if l.startswith(find3):
                    return "char *"
                find4 = " %s :" % member # 位段
                if find4 in l:
                    return l.split(find4)[0].strip()
                find5 = " (*%s)(" % member # 函数
                if find5 in l:
                    return "void *"
            raise InvalidArgsException("%s is not in %s" % (member, structName))

    def showMemberOffset(self, member, structName):
        structName = structName.replace("*", "")
        cmd = "p &((%s*)0)->%s" % (structName, member)
        self.__aCmd.writeLine(cmd)
        s = self.__aCmd.readLine(1)
        if s is None:
            raise InvalidArgsException("%s is not in %s" % (member, structName))
        readStr = s.strip()
        nStr = self.__re1.findall(readStr)[0]
        if "+" not in nStr:
            return 0
        return int(nStr.split("+")[1])

class kprobe(ftrace):
    def __init__(self, args, vm=None, show=False, arch="", stack=False):
        super(kprobe, self).__init__()
        self._vm = vm
        self._gdb = None
        self._c = CexecCmd()
        self._probes = []
        if not show:
            self._checkIsEmpty()
        self._arch = arch
        if self._arch == "":
            self._arch = self._getArchitecture()
        self._reBrackets = re.compile(r"[(](.*?)[)]", re.S)
        self.__args = args
        self._show = show
        self._format = 'x'
        self._stack = stack

    def _getArchitecture(self):
        lines = self._c.cmd('lscpu').split('\n')
        for line in lines:
            if line.startswith("Architecture"):
                arch = line.split(":", 1)[1].strip()
                if arch.startswith("arm"):
                    return "arm"
                if arch.startswith('x86'):
                    return "x86"
                return arch
        return "Unkown"

    def _replaceRegs(self, s):
        if s[-1] in 'UXS':
            self._format = str.lower(s[-1])
            s = s[:-1]
        orig = ['%0', '%1', '%2', '%3', '%4', '%5', '%6', '%7']
        if self._arch == 'x86':
            rep = ['%di', '%si', '%dx', '%cx', '%r8', '%r9', 'bad', 'bad']
        elif self._arch == 'arm':
            rep = ['%r0', '%r1', '%r2', '%r3', 'bad', 'bad']
        elif self._arch == 'aarch64':
            rep = ['%x0', '%x1', '%x2', '%x3', '%x4', '%x5', '%x6', '%x7']
        else:
            raise InvalidArgsException('%s not support', arch)
        for i, o in enumerate(orig):
            if o in s:
                s = s.replace(o, rep[i])
        return s

    def _clearProbes(self):
        if self._show:
            return
        for p in self._probes:
            fPath = self.baseDir + "/tracing/events/kprobes/" + p + "/enable"
            self._echoPath(fPath, 0)
        self._probes = []
        fPath = self.baseDir + "/tracing/kprobe_events"
        self._echoPath(fPath)

    def _splitEgs(self, egs):
        cells = []
        beg = 0
        for i in range(len(egs)):
            if egs[i] == '.':
                cells.append(egs[beg:i])
                cells.append('.')
                beg = i + 1
            elif egs[i] == '-':
                if egs[i + 1] != ">":
                    raise InvalidArgsException("bad %s" % egs)
                cells.append(egs[beg:i])
                cells.append('->')
                beg = i + 2
            elif egs[i] == '>':
                continue
        if beg >= len(egs):
            raise InvalidArgsException("bad %s" % egs)
        cells.append(egs[beg:])
        return cells

    def _skbBodyCheck(self, sMem, type):
        dataTypes = {
            "cdata[": 1,
            "sdata[": 2,
            "ldata[": 4,
            "qdata[": 8,
        }
        memTypes = {
            1 : "__u8",
            2 : "__u16",
            4 : "__u32",
            8 : "__u64",
        }
        for k, v in dataTypes.items():
            if sMem.startswith(k):
                typeOff = self._gdb.showTypeSize(type)
                off = int(sMem[6:-1])
                return [typeOff + off * v, memTypes[v]]
        return [-1, 'None']

    def _skbProc(self, cells, i, cmd):
        layer = cells[i].split(')')[1]
        if layer == "":
            layer = 'l3'
        if layer not in ['l2', 'l3', 'l4']:
            raise InvalidArgsException("%s not support network layser." % layer)
        sHead = self._reBrackets.findall(cells[i])[0]
        sMode = cells[i + 1]
        sMem = cells[i + 2]
        if (sMode != "->"):
            raise InvalidArgsException('$ struct must to use ->')
        off = self._gdb.showMemberOffset('data', 'struct sk_buff')
        cmd = "+0x%x(" % (off) + cmd + ")"

        off = 0
        if sHead == 'struct ethhdr':
            if layer != 'l2':
                InvalidArgsException('can not get eth at %s' % layer)
            off += self._gdb.showMemberOffset(sMem, sHead)
            cells[i + 2] = self._gdb.memberTypes(sMem, sHead)
            return "+0x%x(" % (off) + cmd + ")"
        if layer == 'l2':
            off += self._gdb.showTypeSize('struct ethhdr')
        if sHead == 'struct iphdr':
            if layer == 'l3':
                InvalidArgsException('can not get iphdr at %s' % layer)
            off += self._gdb.showMemberOffset(sMem, sHead)
            cells[i + 2] = self._gdb.memberTypes(sMem, sHead)
            return "+0x%x(" % (off) + cmd + ")"
        if layer != 'l4':
            off += self._gdb.showTypeSize('struct iphdr')
        res = self._skbBodyCheck(sMem, sHead)
        if res[0] >= 0:
            off += res[0]
            cells[i + 2] = res[1]
        else:
            off += self._gdb.showMemberOffset(sMem, sHead)
            cells[i + 2] = self._gdb.memberTypes(sMem, sHead)
        return "+0x%x(" % (off) + cmd + ")"


    def _dispatchStruct(self, cells, i, cmd):
        sStruct = self._reBrackets.findall(cells[i])[0]
        if sStruct in ['struct ethhdr', 'struct iphdr', 'struct icmphdr', 'struct tcphdr', 'struct udphdr']:
            return self._skbProc(cells, i, cmd)
        raise InvalidArgsException('%s not support types' % cells[i])


    def cellsToCmd(self, cells, r):
        i = 0; end = len(cells)
        cmd = r
        while i + 2 < end:
            sHead = cells[i]
            if sHead.startswith("$("):
                cmd = self._dispatchStruct(cells, i, cmd)
            else:
                sMode = cells[i + 1]
                sMem = cells[i + 2]
                if '(' in sHead:
                    sHead = self._reBrackets.findall(sHead)[0]
                off = self._gdb.showMemberOffset(sMem, sHead)
                cells[i + 2] = self._gdb.memberTypes(sMem, sHead)
                if sMode == '.':
                    try:
                        prev, cmd = cmd.split("(", 1)
                    except:
                        raise InvalidArgsException("regs should start with ->.")
                    last = int(prev[1:], 16)
                    cmd = "+0x%x(" % (off + last) + cmd
                elif sMode == '->':
                    cmd = "+0x%x(" % (off) + cmd + ")"
                else:
                    raise InvalidArgsException("%s: bad struct to member mode." % sMode)
            i += 2
        if cells[-1] == "char *":
            cmd +=":string"
        else:
            formDict = {1: 8, 2: 16, 4: 32, 8: 64}
            size = self._gdb.showTypeSize(cells[-1])
            if size in formDict.keys():
                cmd +=":%s%d" % (self._format, formDict[size])
            else:
                raise InvalidArgsException("%s can not show members." % cells[-1])
        return cmd

    "nr=%1~(struct task_struct)->nr_dirtied"
    def transArgs(self, argStr):
        if '=' not in argStr:
            raise InvalidArgsException("no '=' in argStr: %s" % argStr)
        arg, regs = argStr.split('=')
        if arg in ['common_pid', 'common_preempt_count', 'common_flags', 'common_type']:
            raise InvalidArgsException('%s is reserve word, can not used for args' % arg)
        if '~' in regs:
            if self._gdb is None:
                self._gdb = CpaserGdbArgs(self._vm)
            r, egs = regs.split("~")
            cells = self._splitEgs(egs)
            return "%s=%s" %(arg, self.cellsToCmd(cells, self._replaceRegs(r)))
        else:
            return "%s=%s" % (arg, self._replaceRegs(regs))

    def __argSplit(self, arg):
        cells = []
        brackets = 0
        i = beg = 0; end = len(arg)
        while i < end:
            if arg[i] == '(':
                brackets += 1
            elif arg[i] == ')':
                brackets -= 1
            elif brackets == 0 and arg[i] == ' ':
                cells.append(arg[beg:i])
                while arg[i + 1] == ' ':
                    i += 1
                beg = i + 1
            i +=1
        cells.append(arg[beg:])
        if brackets != 0:
            raise InvalidArgsException("%s brackets is not in pair" % arg)
        return cells

    def __transFilter(self, filter, i, beg):
        decisions = ['==', '!=', '~', '>=', '<=', '>', '<']
        s = filter[beg:i]
        for d in decisions:
            if d in s:
                k, v = s.split(d)
                if '_' in k:
                    k, v = invHeadTrans(k, v)
                return "%s%s%s" % (k, d, v)
        raise InvalidArgsException("bad filter format %s" % s)

    def __checkFilter(self, filter):
        cpStr = "()|&"
        i = 0; beg = 0; ret = ""; l = len(filter)
        for i in range(l):
            if filter[i] in cpStr:
                if i and beg != i:
                    ret += self.__transFilter(filter, i, beg)
                beg = i + 1
                ret += filter[i]
        if beg != l:
            ret += self.__transFilter(filter, l, beg)
        return ret

    def __showExpression(self, head, cmd):
        c = cmd.split(' ', 1)[1]
        print("%s %s" % (head, c))

    def __initProbes(self, args):
        i = 0
        if len(args) < 1:
            raise InvalidArgsException("no args.")
        for arg in args:
            self._hex = True
            arg = arg.strip('\n'); cmd = ""
            if len(arg) == 0:
                continue
            a = self.__argSplit(arg.strip())
            head = a[0]
            if head != 'r' and head != 'p':
                raise InvalidArgsException("first arg should be p or r, not %s" % head)
            name = "f%d" % (i); i += 1
            cmd += "%s:%s " % (head, name)
            if '+' in a[1]:
                func = a[1].split('+')[0]
            else:
                func = a[1]
            if not self._checkAvailable(func):
                raise InvalidArgsException("%s is not in available_filter_functions" % func)
            cmd += a[1] + ' '
            if a[-1].startswith("f:"):
                for c in a[2:-1]:
                    cmd += self.transArgs(c) + ' '
                if not self._show:
                    self._echoDPath(self.baseDir + "/tracing/kprobe_events", "'" + cmd + "'")
                try:
                    filter = self.__checkFilter(a[-1].split(':')[1])
                except Exception as e:
                    print(e.message)
                    raise InvalidArgsException('bad filter：%s' % a[-1])
                if self._show:
                    self.__showExpression(head, cmd + 'f:%s' % filter)
                    continue
                else:
                    fPath = self.baseDir + "/tracing/events/kprobes/" + name + "/filter"
                    self._echoPath(fPath, "'" + filter + "'")
            else:
                for c in a[2:]:
                    cmd += self.transArgs(c) + ' '
                if self._show:
                    self.__showExpression(head, cmd)
                    continue
                else:
                    self._echoDPath(self.baseDir + "/tracing/kprobe_events", "'" + cmd + "'")
            fPath = self.baseDir + "/tracing/events/kprobes/" + name + "/enable"
            self._echoPath(fPath, 1)
            self._probes.append(name)
        if self._show:
            return
        if self._stack:
            fPath = self.baseDir + "/tracing/options/stacktrace"
            self._echoPath(fPath, 1)
        else:
            fPath = self.baseDir + "/tracing/options/stacktrace"
            self._echoPath(fPath, 0)

    def _checkIsEmpty(self):
        cmd = 'cat %s/tracing/kprobe_events' % (self.baseDir)
        line = self._c.cmd(cmd).strip()
        if line != "":
            raise FileNotEmptyException("kprobe_events is not empty. should clear other kprobe at first.")

    def _checkAvailable(self, name):
        cmd = "cat " + self.baseDir + "/tracing/available_filter_functions |grep " + name
        ss = self._c.system(cmd).read().strip()
        for res in ss.split('\n'):
            if ':' in res:
                res = res.split(":", 1)[1]
            if ' [' in res:  #for ko symbol
                res = res.split(" [", 1)[0]
            if res == name:
                return True
        return False

    def str2expression(self, s):
        if "->" in s or "." in s:
            pass
        return s

    def procLine(self, line):
        ss = line.split(' ')
        o = ' '
        for s in ss:
            if '=' in s:
                head, value = s.split('=', 1)
                if '_' in head:
                    s = headTrans(head, value)
            o += s + ' '
        print("%s" % o)

    def start(self):
        try:
            self.__initProbes(self.__args)
        except InvalidArgsException as e:
            self._clearProbes()
            del self._gdb
            print(e.message)
            raise InvalidArgsException("input error")
        del self._gdb
        if not self._show:
            super(kprobe, self).start()

    def stop(self):
        self._clearProbes()
        if self._stack:
            fPath = self.baseDir + "/tracing/options/stacktrace"
            self._echoPath(fPath, 0)
        super(kprobe, self).stop()

examples = """examples:"""
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Trace ftrace kprobe events.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=examples
    )
    parser.add_argument('-v', '--vmlinux', type=str, dest='vmlinux', help='set vmlinux path.')
    parser.add_argument('-f', '--file', type=str, dest='file', help='set input args path.')
    parser.add_argument('-a', '--arch', type=str, dest='arch', help='set architecture.')
    parser.add_argument('-s', '--stack', action="store_true", help="show call stacks.")
    parser.add_argument('-S', '--show', action="store_true", help="only show expressions.")
    parser.add_argument(metavar='trace', nargs='+', dest='traces', help='set trace args.')
    args = parser.parse_args()
    traces = args.traces
    if args.file:
        with open(args.file, 'r') as f:
            ts = f.readlines()
            traces += ts
    vm = None
    if args.vmlinux:
        vm = args.vmlinux
    arch = ""
    if args.arch:
        if arch not in ['x86', 'x86_64', 'arm', 'aarch64']:
            raise InvalidArgsException('not support architecture %s' % args.arch)
        arch = args.arch
        if arch.startswith('x86'):
            arch = 'x86'
    k = kprobe(traces, vm, args.show, arch, args.stack)
    k.start()
    if not args.show:
        k.loop()
